import torch
from 数据集 import test_loader, test_ds
from LeNet import LeNet

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = LeNet().to(device)

try:
    state_dict = torch.load('LeNet.pth', map_location=device)
    model.load_state_dict(state_dict)
    print('加载成功')
except:
    print("模型加载失败，无法进行评估")
    exit()


# ---开启评估模式---
loss_fn = torch.nn.NLLLoss()

model.eval()
test_total_loss = 0.0
test_total_true = 0
with torch.no_grad():
    for i, (imgs, labels) in enumerate(test_loader):
        imgs = imgs.to(device)
        labels = labels.to(device)
        # 前向传播、计算损失即可
        y_pre = model(imgs)
        loss = loss_fn(y_pre, labels)
        test_total_loss += loss.item()
        # print(f'step: {i + 1}, loss: {loss.item():.4f}')
        # 统计预测正确的数量
        y_pre_argmax = y_pre.argmax(dim=-1)
        test_total_true += (y_pre_argmax == labels).sum().item()
avd_loss = test_total_loss / len(test_loader)
print(f'测试集平均损失: {avd_loss:.4f}')
print(f'测试集准确率: {test_total_true / len(test_ds):.4f}')